'''
Description: 
Author: Guan Xiongjun
Date: 2022-09-16 13:46:42
LastEditTime: 2022-09-20 12:42:54
'''
import cv2
import torch
import numpy as np
from trans_net import TransNet,TPSGridGen
from torch.autograd import Function, Variable
import itertools
import torch.nn.functional as F

if __name__ == "__main__":
    img = cv2.imread('./C2CL_network/test_data/1.png',0)[np.newaxis,np.newaxis,:,:]
    img=np.float32(img)/255
    image_height, image_width = 480,480
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    img = torch.tensor(img).to(device)
    img = torch.cat([img,img],dim=0)
    n = 4

    net = TransNet(device,image_height, image_width)
    net.to(device=device)
    # x = net(img)
    # net.apply(img,x)


    D = np.zeros((4*4*2,))
    for i in range(len(D)):
        if i%2==1:
            D[i]=0.1

    As = np.array([1,30*np.pi/180,0,0])
    x = np.float32(np.hstack((D,As))[np.newaxis,:])
    x = torch.tensor(x).to(device)
    x = torch.cat([x,x],dim=0)

    img = net.apply(img,x)
    img = torch.squeeze(img).cpu().numpy()[0,:,:]
    img = np.uint8(img*255)
    cv2.imwrite('./C2CL_network/test_data/1_stn.png',img)


